import os
import sys

sys.path.append('.')  # 添加当前目录到 Python 路径
print(os.getcwd())  # 验证下当前工作路径

from tinygraph.graph import TinyGraph
from tinygraph.embedding.open_ai import OpenAIEmb
from tinygraph.llm.open_ai import OpenAILLM

from neo4j import GraphDatabase
from dotenv import load_dotenv  # 用于加载环境变量

# 配置使用的 LLM 和 Embedding 服务，现在只支持 ZhipuAI
# 加载 .env文件, 从而导入api_key
load_dotenv()  # 加载工作目录下的 .env 文件

emb = OpenAIEmb(
    model_name="qwen3-embedding-4b",  # 嵌入模型
    api_key=os.getenv('API_KEY'),
    base_url=os.getenv('BASE_URL')
)
llm = OpenAILLM(
    model_name="qwen3-235b-a22b-instruct",  # LLM 模型
    api_key=os.getenv('API_KEY'),
    base_url=os.getenv('BASE_URL')
)
graph = TinyGraph(
    url="neo4j://localhost:7687",
    username="neo4j",
    password="MySecurePass123!",  # 初次登陆的默认密码为neo4j，此后需修改再使用
    llm=llm,
    emb=emb,
)

# 使用 TinyGraph 添加文档。目前支持所有文本格式的文件。这一步的时间可能较长；
# 结束后，在当前目录下会生成一个 `workspace` 文件夹，包含 `community`、`chunk` 和 `doc` 信息
graph.add_document("example/causal_event_records_202510291600.json")

with graph.driver.session() as session:
    result = session.run("MATCH (n) RETURN count(n) as count")
    count = result.single()["count"]
    print(f"数据库连接正常，节点数量: {count}")

# 执行局部查询测试
local_res = graph.local_query("How much money has been repurchased in total?")
print("\n本地查询结果:")
print(local_res)

global_res = graph.global_query("药明康德有哪些潜在的危机或机遇吗?")
print("\n全局查询结果:")
print(global_res)